# -*- coding: utf-8 -*-
import time
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, \
    QMessageBox, QWidget, QHeaderView, QTableWidgetItem, QAbstractItemView, QPushButton, QMenu
import sys
import os
from PIL import ImageFont
from ultralytics import YOLO

sys.path.append('UIProgram')
from UIProgram.UiMain import Ui_MainWindow
import sys
from PyQt5.QtCore import QTimer, Qt, QThread, pyqtSignal, QCoreApplication
import detect_tools as tools
import cv2
import Config
from UIProgram.QssLoader import QSSLoader
import numpy as np
import torch
import warnings

warnings.filterwarnings('ignore', category=DeprecationWarning)


class MainWindow(QMainWindow):
    def __init__(self, parent=None):
        super(QMainWindow, self).__init__(parent)
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.initMain()
        self.signalconnect()

        # 加载css渲染效果
        style_file = 'UIProgram/style.css'
        qssStyleSheet = QSSLoader.read_qss_file(style_file)
        self.setStyleSheet(qssStyleSheet)

        self.conf = 0.5
        self.unkonwn_text = '无法识别'

        # 添加右键菜单功能
        self.setup_context_menu()

    def setup_context_menu(self):
        # 设置表格支持右键菜单
        self.ui.tableWidget.setContextMenuPolicy(Qt.CustomContextMenu)
        self.ui.tableWidget.customContextMenuRequested.connect(self.show_context_menu)

        # 创建右键菜单
        self.context_menu = QMenu(self)
        self.delete_action = self.context_menu.addAction("删除选中记录")
        self.delete_action.triggered.connect(self.delete_selected_rows)

    def show_context_menu(self, pos):
        # 显示右键菜单
        if self.ui.tableWidget.selectionModel().selectedRows():
            self.context_menu.exec_(self.ui.tableWidget.viewport().mapToGlobal(pos))

    def delete_selected_rows(self):
        # 获取选中的行
        selected_rows = self.ui.tableWidget.selectionModel().selectedRows()
        if not selected_rows:
            return

        # 确认对话框
        reply = QMessageBox.question(
            self, '确认删除',
            '确定要删除选中的记录吗？',
            QMessageBox.Yes | QMessageBox.No,
            QMessageBox.No
        )

        if reply == QMessageBox.Yes:
            # 从后往前删除避免索引问题
            for row in sorted([r.row() for r in selected_rows], reverse=True):
                self.ui.tableWidget.removeRow(row)

    def signalconnect(self):
        self.ui.PicBtn.clicked.connect(self.open_img)
        self.ui.CapBtn.clicked.connect(self.camera_show)
        self.ui.FilesBtn.clicked.connect(self.detact_batch_imgs)
        self.ui.tableWidget.cellClicked.connect(self.on_cell_clicked)

    def initMain(self):
        self.show_width = 770
        self.show_height = 480

        self.org_path = None
        self.current_frame = None
        self.is_camera_open = False
        self.cap = None
        self.capture_btn = None

        self.device = 0 if torch.cuda.is_available() else 'cpu'

        # 加载检测模型
        self.model = YOLO(Config.model_path, task='classify')
        self.model(np.zeros((48, 48, 3)), device=self.device)  # 预先加载推理模型

        # 更新视频图像
        self.timer_camera = QTimer()

        # 表格
        self.ui.tableWidget.verticalHeader().setSectionResizeMode(QHeaderView.Fixed)
        self.ui.tableWidget.verticalHeader().setDefaultSectionSize(40)
        self.ui.tableWidget.setColumnWidth(0, 100)  # 设置列宽
        self.ui.tableWidget.setColumnWidth(1, 300)
        self.ui.tableWidget.setColumnWidth(2, 210)
        self.ui.tableWidget.setColumnWidth(3, 140)
        self.ui.tableWidget.setSelectionBehavior(QAbstractItemView.SelectRows)  # 设置表格整行选中
        self.ui.tableWidget.verticalHeader().setVisible(False)  # 隐藏列标题
        self.ui.tableWidget.setAlternatingRowColors(True)  # 表格背景交替

    def open_img(self):
        if self.cap:
            # 打开图片前关闭摄像头
            self.video_stop()
            self.is_camera_open = False
            self.ui.CapBtn.setText('打开摄像头')
            if self.capture_btn:
                self.capture_btn.hide()
            self.cap = None

        # 弹出的窗口名称：'打开图片'
        # 默认打开的目录：'./'
        file_path, _ = QFileDialog.getOpenFileName(None, '导入影像', './', "Image files (*.jpg *.jpeg *.png *.bmp)")
        if not file_path:
            return

        self.org_path = file_path
        self.org_img = tools.img_cvread(self.org_path)

        # 目标检测
        self.detect_image(self.org_img, file_path)

    def detact_batch_imgs(self):
        # 批量检测图片
        if self.cap:
            # 打开图片前关闭摄像头
            self.video_stop()
            self.is_camera_open = False
            self.ui.CapBtn.setText('打开摄像头')
            if self.capture_btn:
                self.capture_btn.hide()
            self.cap = None
        directory = QFileDialog.getExistingDirectory(self,
                                                     "选取文件夹",
                                                     "./")  # 起始路径
        if not directory:
            return
        # 删除表格所有行
        self.ui.tableWidget.setRowCount(0)
        self.ui.tableWidget.clearContents()
        self.org_path = directory
        img_suffix = ['jpg', 'png', 'jpeg', 'bmp']
        for file_name in os.listdir(directory):
            full_path = os.path.join(directory, file_name)
            if os.path.isfile(full_path) and file_name.split('.')[-1].lower() in img_suffix:
                img_path = full_path
                self.org_img = tools.img_cvread(img_path)
                # 目标检测
                self.detect_image(self.org_img, img_path)
                QApplication.processEvents()  # 刷新页面

    def detect_image(self, img, path=None):
        # 目标检测
        self.results = self.model(img, conf=self.conf)[0]
        self.conf_list = self.results.probs.data.tolist()

        cls_index = self.results.probs.top1
        conf = self.conf_list[cls_index]
        if conf > self.conf:
            show_conf = '{:.2f}%'.format(conf * 100)
            res_name = Config.CH_names[cls_index]
        else:
            show_conf = '100%'
            res_name = self.unkonwn_text

        self.ui.label_conf.setText(show_conf)
        self.ui.resLb.setText(res_name)

        # 显示图片
        now_img = img.copy()
        self.img_width, self.img_height = self.get_resize_size(now_img)
        resize_cvimg = cv2.resize(now_img, (self.img_width, self.img_height))
        pix_img = tools.cvimg_to_qpiximg(resize_cvimg)
        self.ui.label_show.setPixmap(pix_img)
        self.ui.label_show.setAlignment(Qt.AlignCenter)

        # 添加到表格
        self.tabel_info_show(res_name, show_conf, path=path)
        self.ui.tableWidget.scrollToBottom()

    def video_start(self):
        # 删除表格所有行
        self.ui.tableWidget.setRowCount(0)
        self.ui.tableWidget.clearContents()

        # 定时器开启，每隔一段时间，读取一帧
        self.timer_camera.start(30)  # 约30ms一帧
        self.timer_camera.timeout.connect(self.open_frame)

    def tabel_info_show(self, cls, conf, path=None):
        row_count = self.ui.tableWidget.rowCount()
        self.ui.tableWidget.insertRow(row_count)

        # 标准化路径：统一为系统分隔符，并转为绝对路径
        safe_path = os.path.abspath(os.path.normpath(path)) if path else ""
        # 获取文件名
        file_name = os.path.basename(safe_path) if path else "Camera" if self.is_camera_open else "Video Frame"

        # 存储到表格（完整路径作为隐藏数据，显示文件名）
        self.ui.tableWidget.setItem(row_count, 0, QTableWidgetItem(str(row_count + 1)))
        # 第2列显示文件名，但存储完整路径在Qt.UserRole中
        path_item = QTableWidgetItem(file_name)
        path_item.setData(Qt.UserRole, safe_path)  # 存储完整路径
        self.ui.tableWidget.setItem(row_count, 1, path_item)
        self.ui.tableWidget.setItem(row_count, 2, QTableWidgetItem(str(cls)))
        self.ui.tableWidget.setItem(row_count, 3, QTableWidgetItem(str(conf)))

    def video_stop(self):
        if self.cap:
            self.cap.release()
        self.timer_camera.stop()

    def open_frame(self):
        ret, now_img = self.cap.read()
        if ret:
            self.current_frame = now_img.copy()  # 保存当前帧

            # 获取缩放后的图片尺寸
            self.img_width, self.img_height = self.get_resize_size(now_img)
            resize_cvimg = cv2.resize(now_img, (self.img_width, self.img_height))
            pix_img = tools.cvimg_to_qpiximg(resize_cvimg)
            self.ui.label_show.setPixmap(pix_img)
            self.ui.label_show.setAlignment(Qt.AlignCenter)
        else:
            self.video_stop()

    def camera_show(self):
        if not self.is_camera_open:
            # 打开摄像头
            self.ui.CapBtn.setText('关闭摄像头')
            self.cap = cv2.VideoCapture(0)

            # 创建拍照按钮
            if not self.capture_btn:
                self.capture_btn = QPushButton('拍照', self)
                self.capture_btn.clicked.connect(self.capture_image)
                # 设置按钮样式和位置
                self.capture_btn.setStyleSheet("""
                    QPushButton {
                        background-color: #4CAF50;
                        color: white;
                        border: none;
                        padding: 8px 16px;
                        font-size: 14px;
                        border-radius: 4px;
                    }
                    QPushButton:hover {
                        background-color: #45a049;
                    }
                """)
                self.capture_btn.move(20, 20)  # 左上角位置

            self.capture_btn.show()
            self.is_camera_open = True
            self.video_start()
        else:
            # 关闭摄像头
            self.ui.CapBtn.setText('打开摄像头')
            if self.capture_btn:
                self.capture_btn.hide()
            if self.cap:
                self.cap.release()
            self.is_camera_open = False
            self.ui.label_show.clear()

    def capture_image(self):
        if not self.is_camera_open or self.current_frame is None:
            return

        # 保存图片
        save_path = os.path.join(os.getcwd(), 'captured_images')
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        timestamp = time.strftime("%Y%m%d_%H%M%S")
        file_path = os.path.join(save_path, f"capture_{timestamp}.jpg")
        cv2.imwrite(file_path, self.current_frame)

        # 显示保存成功提示
        QMessageBox.information(self, "拍照成功", f"照片已保存到:\n{file_path}")

        # 进行目标检测并显示结果
        self.detect_image(self.current_frame, file_path)

    def get_resize_size(self, img):
        _img = img.copy()
        img_height, img_width, depth = _img.shape
        ratio = img_width / img_height
        if ratio >= self.show_width / self.show_height:
            self.img_width = self.show_width
            self.img_height = int(self.img_width / ratio)
        else:
            self.img_height = self.show_height
            self.img_width = int(self.img_height * ratio)
        return self.img_width, self.img_height

    def on_cell_clicked(self, row, column):
        try:
            # 1. 获取表格中的路径项
            item_path = self.ui.tableWidget.item(row, 1)
            if not item_path:
                QMessageBox.warning(self, "错误", "该记录无有效路径")
                return

            # 2. 从UserRole获取完整路径
            img_path = item_path.data(Qt.UserRole)
            if not img_path:
                QMessageBox.warning(self, "错误", "该记录无有效路径")
                return

            # 3. 验证路径有效性（Camera和Video Frame特殊处理）
            if img_path in ["Camera", "Video Frame"]:
                return

            img_path = os.path.abspath(os.path.normpath(img_path))  # 再次标准化确保安全
            if not os.path.exists(img_path):
                display_path = img_path.replace("\\", "/")
                QMessageBox.critical(
                    self,
                    "文件不存在",
                    f"无法找到图片：\n{display_path}\n"
                    "可能原因：\n"
                    "1. 文件被移动/删除\n"
                    "2. 路径名包含特殊字符\n"
                    "3. 程序无访问权限"
                )
                return

            # 4. 加载图片
            now_img = tools.img_cvread(img_path)
            if now_img is None:
                raise ValueError("图片解码失败")

            # 5. 显示图片
            self.img_width, self.img_height = self.get_resize_size(now_img)
            resize_cvimg = cv2.resize(now_img, (self.img_width, self.img_height))
            pix_img = tools.cvimg_to_qpiximg(resize_cvimg)

            # 释放旧资源（避免内存泄漏）
            if hasattr(self, 'current_pixmap'):
                self.current_pixmap = None
            self.ui.label_show.setPixmap(pix_img)
            self.current_pixmap = pix_img  # 保留引用

            # 6. 更新文字信息
            cls_item = self.ui.tableWidget.item(row, 2)
            conf_item = self.ui.tableWidget.item(row, 3)
            if cls_item and conf_item:
                self.ui.resLb.setText(cls_item.text())
                self.ui.label_conf.setText(conf_item.text())

        except Exception as e:
            QMessageBox.critical(self, "错误", f"加载失败：{str(e)}")
            print(f"[ERROR] {e}")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = MainWindow()
    win.show()
    sys.exit(app.exec_())